from typing import List

import langchain_core
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage

from ai_engine.core.model.chat import ChatMessage


class ParamUtil:
    _ROLE_MAP = {"human": "用户: ", "ai": "Assistant: "}

    @staticmethod
    def convert_to_message_list(chat_history: List[ChatMessage]) -> List[BaseMessage]:
        """Convert ChatMessage to a LangChain message.
        """
        return [ParamUtil.convert_to_message(msg) for msg in chat_history]

    @staticmethod
    def convert_to_message(message: ChatMessage) -> BaseMessage:
        """Convert ChatMessage to a LangChain message.
        """
        role = message.role
        if role == "human":
            return HumanMessage(content=message.content)
        elif role == "assistant":
            return AIMessage(content=message.content)
        elif role == "system":
            return SystemMessage(content=message.content)
        else:
            return langchain_core.messages.ChatMessage(content=message.content, role=role)

    @staticmethod
    def get_chat_history_str(chat_history: List[ChatMessage]) -> str:
        buffer = ""
        for msg in chat_history:
            role_prefix = ParamUtil._ROLE_MAP.get(msg.role, f"{msg.role}: ")
            buffer += f"\n{role_prefix}{msg.content}"
        return buffer
